/*
 * MultirootSolverImp.cpp
 *
 *  Created on: 23 Jun 2011
 *      Author: allan
 */

#include "MultirootSolver.h"

// C++ includes
#include <iostream>
using namespace std;

// GSL includes
#include <gsl/gsl_multiroots.h>

class MultirootSolverImp
{
public:
	MultirootSolverImp();
	
	MultirootSolverImp(unsigned dimension, MultirootAlgorithm algorithm);
	
	~MultirootSolverImp();
	
	void Initialiaze(unsigned dimension, MultirootAlgorithm algorithm);
	
	void SetSolverOptions(double tolerance, unsigned maxiterations);
	
	bool Iterate(MultirootProblem& problem, VectorXd& x);
	
	bool Solve(MultirootProblem& problem, VectorXd& x);
	
public:
	struct Params
	{
		MultirootProblem* problem;
		VectorXd x;
		VectorXd F;
		MatrixXd J;
	};
	
private:
	unsigned dimension;
	
	Params params;
	
	double tolerance;

	unsigned maxiterations;
	
	gsl_multiroot_function_fdf func;
	
	const gsl_multiroot_fdfsolver_type* type;
	
	gsl_multiroot_fdfsolver* solver;
	
	gsl_vector* xGSL;
};

MultirootSolver::MultirootSolver()
{
	imp = std::unique_ptr<MultirootSolverImp>(new MultirootSolverImp());
}

MultirootSolver::MultirootSolver(unsigned dimension, MultirootAlgorithm algorithm)
{
	imp = std::unique_ptr<MultirootSolverImp>(new MultirootSolverImp(dimension, algorithm));
}

MultirootSolver::~MultirootSolver()
{}

void MultirootSolver::Initialiaze(unsigned dimension, MultirootAlgorithm algorithm)
{
	imp->Initialiaze(dimension, algorithm);
}

void MultirootSolver::SetSolverOptions(double tolerance, unsigned maxiterations)
{
	imp->SetSolverOptions(tolerance, maxiterations);
}

bool MultirootSolver::Iterate(MultirootProblem& problem, VectorXd& x)
{
	return imp->Iterate(problem, x);
}

bool MultirootSolver::Solve(MultirootProblem& problem, VectorXd& x)
{
	return imp->Solve(problem, x);
}

int GSLFunction(const gsl_vector* x, void* params, gsl_vector* F);
int GSLJacobian(const gsl_vector* x, void* params, gsl_matrix* J);
int GSLFunctionJacobian(const gsl_vector* x, void* params, gsl_vector* F, gsl_matrix* J);

MultirootSolverImp::MultirootSolverImp() :
dimension(0), tolerance(1.0E-6), maxiterations(100), 
type(gsl_multiroot_fdfsolver_gnewton), solver(0), xGSL(0)
{}

MultirootSolverImp::MultirootSolverImp(unsigned dimension, MultirootAlgorithm algorithm) :
dimension(0), tolerance(1.0E-6), maxiterations(100), 
type(gsl_multiroot_fdfsolver_gnewton), solver(0), xGSL(0)
{
	Initialiaze(dimension, algorithm);
}

MultirootSolverImp::~MultirootSolverImp()
{
	gsl_multiroot_fdfsolver_free(solver);
	gsl_vector_free(xGSL);
}

void MultirootSolverImp::Initialiaze(unsigned dimension, MultirootAlgorithm algorithm)
{
	// Set the data member (dimension)
	this->dimension = dimension;
	
	// Detemine the (type) of the GSL multiroot algorithm 
	switch(algorithm)
	{
	case Hybridj:
		type = gsl_multiroot_fdfsolver_hybridj;
		break;
	case Hybridsj:
		type = gsl_multiroot_fdfsolver_hybridsj;
		break;
	case StandardNewton:
		type = gsl_multiroot_fdfsolver_newton;
		break;
	case ModifiedNewton:
		type = gsl_multiroot_fdfsolver_gnewton;
		break;
	default:
		type = gsl_multiroot_fdfsolver_gnewton;
		break;
	}
	
	// Set the data member (params)
	params.problem = 0; // This field is set in the methods Iterate and Solve
	params.x       = VectorXd::Zero(dimension);
	params.F       = VectorXd::Zero(dimension);
	params.J       = MatrixXd::Zero(dimension, dimension);
	
	// Free the GSL objects (solver) and (xGSL) if they have already been initialized
	if(solver != 0) gsl_multiroot_fdfsolver_free(solver);
	if(xGSL   != 0) gsl_vector_free(xGSL);

	// Initialiaze the GSL objects (solver) and (xGSL)
	solver = gsl_multiroot_fdfsolver_alloc(type, dimension);
	xGSL   = gsl_vector_alloc(dimension);
	
	// Initialiaze the GSL object (func)
	func.n      = dimension;
	func.f      = &GSLFunction;
	func.df     = &GSLJacobian;
	func.fdf    = &GSLFunctionJacobian;
	func.params = &params;
}

void MultirootSolverImp::SetSolverOptions(double tolerance, unsigned maxiterations)
{
	this->tolerance = tolerance;
	this->maxiterations = maxiterations;
}

bool MultirootSolverImp::Iterate(MultirootProblem& problem, VectorXd& x)
{
	// Set the (problem) field of the instance (params) 
	params.problem = &problem;
	
	// Transfer the data in the Eigen vector (x) to the GSL vector (xGSL)
	for(unsigned i = 0; i < dimension; ++i) gsl_vector_set(xGSL, i, x[i]);
	
	// Set the GSL object (solver) with the function structure (func) and the vector (xGSL)
	gsl_multiroot_fdfsolver_set(solver, &func, xGSL);
	
	// Perform one iteration
	gsl_multiroot_fdfsolver_iterate(solver);

	// Transfer the data in the GSL vector (xGSL) to the Eigen vector (x)
	for(unsigned i = 0; i < dimension; ++i) x[i] = gsl_vector_get(solver->x, i);
	
	return (gsl_multiroot_test_residual(solver->f, tolerance) == GSL_CONTINUE);
}

bool MultirootSolverImp::Solve(MultirootProblem& problem, VectorXd& x)
{
	// Set the (problem) field of the instance (params) 
	params.problem = &problem;
	
	// An axiliary variable to indicate errors during the iterations 
	int status = GSL_CONTINUE;
	
	// An axiliary variable to indicate if the solver got stuck
	bool stuck = false;
	
	// Transfer the data in the Eigen vector (x) to the GSL vector (xGSL)
	for(unsigned i = 0; i < dimension; ++i) gsl_vector_set(xGSL, i, x[i]);

	// Set the GSL object (solver) with the function structure (func) and the vector (xGSL)
	gsl_multiroot_fdfsolver_set(solver, &func, xGSL);
	
	unsigned iteration = 0;
	
	while(status == GSL_CONTINUE && iteration < maxiterations)
	{
		// Perform one iteration and get the status
		status = gsl_multiroot_fdfsolver_iterate(solver);
		
		// Check if the solver is stuck
		if(status) { stuck = true; break; }
		
		// Transfer the data in the GSL vector (xGSL) to the Eigen vector x
		for(unsigned i = 0; i < dimension; ++i) x[i] = gsl_vector_get(solver->x, i);
	
		// Check the residual of the solution in terms of the function values
		status = gsl_multiroot_test_residual(solver->f, tolerance);
		
		// Increment the data member (iteration)
		++iteration;
	}
	
	// Transfer the data in the GSL vector (xGSL) to the Eigen vector (x)
	for(unsigned i = 0; i < dimension; ++i) x[i] = gsl_vector_get(solver->x, i);
	
	// Check if the method converged
	bool converged = iteration < maxiterations && !stuck;
	
	return converged;
}

int GSLFunction(const gsl_vector* gslx, void* params, gsl_vector* gslF)
{
	// Cast (params) to the correct type
	MultirootSolverImp::Params& p = *((MultirootSolverImp::Params*)params);
	
	// Transfer the data in the GSL vector (gslx) to the Eigen vector (p.x)
	for(unsigned i = 0; i < gslx->size; ++i) p.x[i] = gsl_vector_get(gslx, i);
	
	// Compute (F) using the multiroot problem object
	p.problem->Function(p.x, p.F);
	
	// Transfer the data in the Eigen vector (p.F) to the GSL vector (gslF)
	for(unsigned i = 0; i < gslx->size; ++i) gsl_vector_set(gslF, i, p.F[i]);
	
	return GSL_SUCCESS;
}

int GSLJacobian(const gsl_vector* gslx, void* params, gsl_matrix* gslJ)
{
	// Cast (params) to the correct type
	MultirootSolverImp::Params& p = *((MultirootSolverImp::Params*)params);
	
	// Transfer the data in the GSL vector (gslx) to the Eigen vector (p.x)
	for(unsigned i = 0; i < gslx->size; ++i) p.x[i] = gsl_vector_get(gslx, i);
	
	// Compute (J) using the multiroot problem object
	p.problem->Jacobian(p.x, p.J);
	
	// Transfer the data in the Eigen matrix (J) to the GSL matrix (gslJ) 
	for(unsigned i = 0; i < gslx->size; ++i) for(unsigned j = 0; j < gslx->size; ++j)
		gsl_matrix_set(gslJ, i, j, p.J(i, j));
	
	return GSL_SUCCESS;
}

int GSLFunctionJacobian(const gsl_vector* gslx, void* params, gsl_vector* gslF, gsl_matrix* gslJ)
{
	// Cast (params) to the correct type
	MultirootSolverImp::Params& p = *((MultirootSolverImp::Params*)params);
	
	// Transfer the data in the GSL vector (gslx) to the Eigen vector (p.x)
	for(unsigned i = 0; i < gslx->size; ++i) p.x[i] = gsl_vector_get(gslx, i);
	
	// Compute (F) and (J) using the multiroot problem object
	p.problem->FunctionJacobian(p.x, p.F, p.J);
	
	// Transfer the data in the Eigen vector (p.F) to the GSL vector (gslF)
	for(unsigned i = 0; i < gslx->size; ++i) gsl_vector_set(gslF, i, p.F[i]);
	
	// Transfer the data in the Eigen matrix (J) to the GSL matrix (gslJ) 
	for(unsigned i = 0; i < gslx->size; ++i) for(unsigned j = 0; j < gslx->size; ++j)
		gsl_matrix_set(gslJ, i, j, p.J(i, j));
	
	return GSL_SUCCESS;
}
